{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualization - Custom Components\n",
    "\n",
    "### The Boltzmann Wealth Model "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you want to get straight to the tutorial checkout these environment providers:<br>\n",
    "[![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/projectmesa/mesa/main?labpath=docs%2Ftutorials%2F4_visualization_basic.ipynb) (This can take 30 seconds to 5 minutes to load)\n",
    "\n",
    "Due to conflict with Colab and Solara there are no colab links for this tutorial\n",
    "\n",
    "*If you are running locally, please ensure you have the latest Mesa version installed.*\n",
    "\n",
    "## Tutorial Description\n",
    "\n",
    "This tutorial extends the Boltzmann wealth model from the [Visualization Basic Dashboard tutorial](https://mesa.readthedocs.io/latest/tutorials/4_visualization_basic.html), by adding an interactive dashboard. \n",
    "\n",
    "In this portion, we will demonstrate how users can employ create dynamic agent representation with their Mesa dashboards. This is part two of three visualization tutorials. \n",
    "\n",
    "*If you are starting here please see the [Running Your First Model tutorial](https://mesa.readthedocs.io/latest/tutorials/0_first_model.html) for dependency and start-up instructions*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import Dependencies\n",
    "This includes importing of dependencies needed for the tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# Has multi-dimensional arrays and matrices.\n",
    "# Has a large collection of mathematical functions to operate on these arrays.\n",
    "import numpy as np\n",
    "\n",
    "# Data manipulation and analysis.\n",
    "import pandas as pd\n",
    "\n",
    "# Data visualization tools.\n",
    "import seaborn as sns\n",
    "\n",
    "import mesa\n",
    "from mesa.discrete_space import CellAgent, OrthogonalMooreGrid\n",
    "from mesa.visualization import SolaraViz, make_plot_component, make_space_component"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Basic Model\n",
    "\n",
    "The following is the basic model we will be using to build the dashboard. This is the same model seen in tutorials 0-3. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_gini(model):\n",
    "    agent_wealths = [agent.wealth for agent in model.agents]\n",
    "    x = sorted(agent_wealths)\n",
    "    N = model.num_agents\n",
    "    B = sum(xi * (N - i) for i, xi in enumerate(x)) / (N * sum(x))\n",
    "    return 1 + (1 / N) - 2 * B\n",
    "\n",
    "\n",
    "class MoneyAgent(CellAgent):\n",
    "    \"\"\"An agent with fixed initial wealth.\"\"\"\n",
    "\n",
    "    def __init__(self, model, cell):\n",
    "        \"\"\"initialize a MoneyAgent instance.\n",
    "\n",
    "        Args:\n",
    "            model: A model instance\n",
    "        \"\"\"\n",
    "        super().__init__(model)\n",
    "        self.cell = cell\n",
    "        self.wealth = 1\n",
    "\n",
    "    def move(self):\n",
    "        \"\"\"Move the agent to a random neighboring cell.\"\"\"\n",
    "        self.cell = self.cell.neighborhood.select_random_cell()\n",
    "\n",
    "    def give_money(self):\n",
    "        \"\"\"Give 1 unit of wealth to a random agent in the same cell.\"\"\"\n",
    "        cellmates = [a for a in self.cell.agents if a is not self]\n",
    "\n",
    "        if cellmates:  # Only give money if there are other agents present\n",
    "            other = self.random.choice(cellmates)\n",
    "            other.wealth += 1\n",
    "            self.wealth -= 1\n",
    "\n",
    "    def step(self):\n",
    "        \"\"\"do one step of the agent.\"\"\"\n",
    "        self.move()\n",
    "        if self.wealth > 0:\n",
    "            self.give_money()\n",
    "\n",
    "\n",
    "class MoneyModel(mesa.Model):\n",
    "    \"\"\"A model with some number of agents.\"\"\"\n",
    "\n",
    "    def __init__(self, n=10, width=10, height=10, seed=None):\n",
    "        \"\"\"Initialize a MoneyModel instance.\n",
    "\n",
    "        Args:\n",
    "            N: The number of agents.\n",
    "            width: width of the grid.\n",
    "            height: Height of the grid.\n",
    "        \"\"\"\n",
    "        super().__init__(seed=seed)\n",
    "        self.num_agents = n\n",
    "        self.grid = OrthogonalMooreGrid((width, height), random=self.random)\n",
    "\n",
    "        # Create agents\n",
    "        MoneyAgent.create_agents(\n",
    "            self,\n",
    "            self.num_agents,\n",
    "            self.random.choices(self.grid.all_cells.cells, k=self.num_agents),\n",
    "        )\n",
    "\n",
    "        self.datacollector = mesa.DataCollector(\n",
    "            model_reporters={\"Gini\": compute_gini}, agent_reporters={\"Wealth\": \"wealth\"}\n",
    "        )\n",
    "        self.datacollector.collect(self)\n",
    "\n",
    "    def step(self):\n",
    "        \"\"\"do one step of the model\"\"\"\n",
    "        self.agents.shuffle_do(\"step\")\n",
    "        self.datacollector.collect(self)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Lets make sure the model works\n",
    "model = MoneyModel(100, 10, 10)\n",
    "for _ in range(20):\n",
    "    model.step()\n",
    "\n",
    "\n",
    "data = model.datacollector.get_agent_vars_dataframe()\n",
    "# Use seaborn\n",
    "g = sns.histplot(data[\"Wealth\"], discrete=True)\n",
    "g.set(title=\"Wealth distribution\", xlabel=\"Wealth\", ylabel=\"Number of Agents\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Adding visualization\n",
    "\n",
    "So far, we've built a model, run it, and analyzed some output afterwards. However, one of the advantages of agent-based models is that we can often watch them run step by step, potentially spotting unexpected patterns, behaviors or bugs, or developing new intuitions, hypotheses, or insights. Other times, watching a model run can explain it to an unfamiliar audience better than static explanations. Like many ABM frameworks, Mesa allows you to create an interactive visualization of the model. In this section we'll walk through creating a visualization using built-in components, and (for advanced users) how to create a new visualization element.\n",
    "\n",
    "First, a quick explanation of how Mesa's interactive visualization works. The visualization is done in a browser window or Jupyter instance, using the [Solara](https://solara.dev/) framework, a pure Python, React-style web framework. Running `solara run app.py` will launch a web server, which runs the model, and displays model detail at each step via a plotting library. Alternatively, you can execute everything inside a Jupyter instance and display it inline.\n",
    "\n",
    "*Thanks to @Corvince for all his work creating Mesa's visualization capability*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building Custom Components\n",
    "\n",
    "This section is for users who have a basic familiarity with Python's Matplotlib plotting library.\n",
    "\n",
    "If the visualization elements provided by Mesa aren't enough for you, you can build your own and plug them into the model server.\n",
    "\n",
    "For this example, let's build a simple histogram visualization, which can count the number of agents with each value of wealth."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First we need to update our imports\n",
    "\n",
    "We use Matplotlib in this tutorial, but Mesa also has Altair. If you would like other visualization support like Plotly or Bokeh, please feel free to [contribute](https://github.com/projectmesa/mesa/blob/main/CONTRIBUTING.md)\n",
    "\n",
    "In addition, due to the way Solara works we need to trigger an update whenever the underlying model changes. For this you need to register an update counter with every component."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import solara\n",
    "from matplotlib.figure import Figure\n",
    "\n",
    "from mesa.visualization.utils import update_counter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we provide a function for our agent portrayal and our model parameters. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def agent_portrayal(agent):\n",
    "    size = 10\n",
    "    color = \"tab:red\"\n",
    "    if agent.wealth > 0:\n",
    "        size = 50\n",
    "        color = \"tab:blue\"\n",
    "    return {\"size\": size, \"color\": color}\n",
    "\n",
    "\n",
    "model_params = {\n",
    "    \"n\": {\n",
    "        \"type\": \"SliderInt\",\n",
    "        \"value\": 50,\n",
    "        \"label\": \"Number of agents:\",\n",
    "        \"min\": 10,\n",
    "        \"max\": 100,\n",
    "        \"step\": 1,\n",
    "    },\n",
    "    \"width\": 10,\n",
    "    \"height\": 10,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Page Tab View\n",
    "\n",
    "#### **Plot Components**\n",
    "   \n",
    "You can place different components (except the renderer) on separate pages according to your preference. There are no restrictions on page numbering — pages do not need to be sequential or positive. Each page acts as an independent window where components may or may not exist.\n",
    "\n",
    "The default page is `page=0`. If pages are not sequential (e.g., `page=1` and `page=10`), the system will automatically create the 8 empty pages in between to maintain consistent indexing. To avoid empty pages in your dashboard, use sequential page numbers.\n",
    "\n",
    "To assign a plot component to a specific page, pass the `page` keyword argument to `make_plot_component`. For example, the following will display the plot component on page 1:\n",
    "\n",
    "```python\n",
    "plot_comp = make_plot_component(\"encoding\", page=1)\n",
    "```\n",
    "\n",
    "#### **Custom Components**\n",
    "If you want a custom component to appear on a specific page, you must pass it as a tuple containing the component and the page number.\n",
    "\n",
    "```python\n",
    "@solara.component\n",
    "def CustomComponent():\n",
    "    ...\n",
    "\n",
    "page = SolaraViz(\n",
    "    model,\n",
    "    renderer,\n",
    "    components=[(CustomComponent, 1)]  # Custom component will appear on page 1\n",
    ")\n",
    "```\n",
    "\n",
    "> ⚠️ **Warning**\n",
    "> Running the model can be performance-intensive. It is strongly recommended to pause the model in the dashboard before switching pages."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we add our custom component. In this case we will build a histogram of agent wealth. \n",
    "\n",
    "Besides the standard matplotlib code to build a histogram, please notice 3 key features. \n",
    "\n",
    "1. `@solara.component` this is needed for any compoenent you add\n",
    "2. `update_counter.get()` this is needed so solara updates the dashboard with your agent based model \n",
    "3. you must initialize a `figure` using this method instead of `plt.figure()`, for thread safety purpose\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@solara.component\n",
    "def Histogram(model):\n",
    "    update_counter.get()  # This is required to update the counter\n",
    "    # Note: you must initialize a figure using this method instead of\n",
    "    # plt.figure(), for thread safety purpose\n",
    "    fig = Figure()\n",
    "    ax = fig.subplots()\n",
    "    wealth_vals = [agent.wealth for agent in model.agents]\n",
    "    # Note: you have to use Matplotlib's OOP API instead of plt.hist\n",
    "    # because plt.hist is not thread-safe.\n",
    "    ax.hist(wealth_vals, bins=10)\n",
    "    solara.FigureMatplotlib(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we create the model an initialize the visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create initial model instance\n",
    "money_model = MoneyModel(n=50, width=10, height=10)\n",
    "\n",
    "SpaceGraph = make_space_component(agent_portrayal)\n",
    "GiniPlot = make_plot_component(\"Gini\", page=1)\n",
    "\n",
    "page = SolaraViz(\n",
    "    money_model,\n",
    "    components=[SpaceGraph, GiniPlot, (Histogram, 2)],\n",
    "    model_params=model_params,\n",
    "    name=\"Boltzmann Wealth Model\",\n",
    ")\n",
    "# This is required to render the visualization in the Jupyter notebook\n",
    "page"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can even run the visuals independently by calling it with the model instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Histogram(money_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise\n",
    " - Build you own custom component"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Next Steps\n",
    "\n",
    "Check out the next [batch run tutorial](https://mesa.readthedocs.io/latest/tutorials/9_batch_run.html) on how to conduct parameter sweeps and run numerous iterations of your model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[Comer2014] Comer, Kenneth W. “Who Goes First? An Examination of the Impact of Activation on Outcome Behavior in AgentBased Models.” George Mason University, 2014. http://mars.gmu.edu/bitstream/handle/1920/9070/Comer_gmu_0883E_10539.pdf\n",
    "\n",
    "[Dragulescu2002] Drăgulescu, Adrian A., and Victor M. Yakovenko. “Statistical Mechanics of Money, Income, and Wealth: A Short Survey.” arXiv Preprint Cond-mat/0211175, 2002. http://arxiv.org/abs/cond-mat/0211175."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.5"
  },
  "widgets": {
   "state": {},
   "version": "1.1.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
